##########################################################################################

library(data.table)
library(optparse)
library(ArchR)
library(dplyr)

##########################################################################################
option_list <- list(
    make_option(c("--comine_data_file"), type = "character"),
    make_option(c("--mean_expr_file"), type = "character"),
    make_option(c("--cluster"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    ## 整合atac和rna的文件
    comine_data_file <- "~/20231121_singleMuti/results/qc_atac_v3/germ/testis_combined_peak.combineRNA.qc.Rdata"

    ## cluster
    cluster <- "germ"

    ## 表达文件
    mean_expr_file <- "~/20231121_singleMuti/results/qc_atac_v3/germ/GeneExpression.MeanByCellType.tsv"

    ## 输出
    out_path <- "~/20231121_singleMuti/results/celltype_plot/exp_score"

}

###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

comine_data_file <- opt$comine_data_file
cluster <- opt$cluster
mean_expr_file <- opt$mean_expr_file
out_path <- opt$out_path

dir.create(out_path , recursive = T)

###########################################################################################
## 导入数据
a <- load(comine_data_file)
## atac_proj
if( cluster %in% c("all" , "somatic" , "germ") ){
    atac_proj <- testis_combined_peak_combineRNA
    #atac_proj <- addPeak2GeneLinks(ArchRProj = atac_proj , useMatrix = "GeneExpressionMatrix")
}

mean_expr <- fread(mean_expr_file)

###########################################################################################
## 计算基因表达和开放的相关系数
corGSM_GIM <- correlateMatrices(
    ArchRProj = atac_proj,
    useMatrix1 = "GeneScoreMatrix",
    useMatrix2 = "GeneExpressionMatrix",
    reducedDims = "IterativeLSI"
)

## 提取基因名在两者间完全一致的
index <- which(corGSM_GIM$GeneScoreMatrix_name == corGSM_GIM$GeneExpressionMatrix_name)
corGSM_GIM <- corGSM_GIM[index,]

## 去除无法计算相关系数的基因
corGSM_GIM <- subset(corGSM_GIM , cor!="NA")

## 重新矫正p值
corGSM_GIM$fdr <- p.adjust( corGSM_GIM$pval , method = "fdr" )

###########################################################################################
## 提取peak-gene的对应关系
corrCutoff <- 0.5

p2gMat <- plotPeak2GeneHeatmap(
  atac_proj, 
  corCutOff = corrCutoff, 
  groupBy="cell_type",
  nPlot = 1000000, 
  returnMatrices=TRUE, 
  k=25, seed=1)

## 输出peak-gene相关参数
kclust_df_out <- cbind( p2gMat$Peak2GeneLinks )
tmp_peak_dat <- data.frame(atac_proj@peakSet)
tmp_peak_dat$peak <- paste0( tmp_peak_dat$seqnames , ":" , tmp_peak_dat$start , "-" , tmp_peak_dat$end )
kclust_df_out <- merge( kclust_df_out , tmp_peak_dat , by = "peak" )

out_file <- paste0( out_path , "/" , cluster , "_peak2gene.tsv" )
write.table(kclust_df_out , out_file , row.names = F , sep = "\t")

## 计算每个基因连接多少peak
data_peakNum <- data.frame(kclust_df_out) %>%
group_by( gene ) %>%
summarize( peak_num = length(unique(peak)) )

## 提取存在表达的基因
mean_expr <- mean_expr[apply(mean_expr[,-1] , 1 , max) > 0,]
expr_peak <- merge( mean_expr , data_peakNum , all.x = T )
expr_peak$peak_num[is.na(expr_peak$peak_num)] <- 0

out_file <- paste0( out_path , "/" , cluster , "_meanExp_linkPeakNum.tsv" )
write.table(expr_peak , out_file , row.names = F , sep = "\t")

###########################################################################################
## 连接相关系数以及peak-gene的关系
corGSM_GIM$gene <- corGSM_GIM$GeneExpressionMatrix_name
result <- merge( corGSM_GIM , data_peakNum , all.x = T )
result <- result[,c("gene" , "cor" , "fdr" , "peak_num")]
colnames(result) <- c("gene" , "cor_exp2atac" , "fdr_exp2atac" , "linkpeak_num")
result$linkpeak_num[is.na(result$linkpeak_num)] <- 0

out_file <- paste0( out_path , "/" , cluster , "_exp-atac_linkPeakNum.tsv" )
write.table(result , out_file , row.names = F , sep = "\t")